package top.lingma.algorithm;

import lombok.Data;
import org.aspectj.lang.JoinPoint;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.integration.support.locks.LockRegistry;
import top.lingma.annotation.RateLimit;
import top.lingma.processor.insufficient.InsufficientLimitProcessor;

import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Lock;

/**
 * 令牌桶限流器
 *
 * @author yui
 * @date 2019年05月21日 上午10:46
 */
@Data
public class RedisRateLimiter {
    InsufficientLimitProcessor tokenInsufficient;
    JoinPoint joinPoint;
    RateLimit limit;

    /**
     * redis key
     */
    private String key;
    /**
     * redis分布式锁的key
     *
     * @return
     */
    private String lockKey;
    Long permitsPerSecond;
    Integer maxBurstSeconds;

    /**
     * 分布式同步锁
     */
    private LockRegistry syncLock;

    private RedisTemplate redisTemplate;


    public RedisRateLimiter(String key, Long permitsPerSecond, Integer maxBurstSeconds, RedisTemplate redisTemplate, LockRegistry syncLock) {
        this.key = key;
        this.lockKey = "LIMITER_LOCK:" + key;
        this.permitsPerSecond = permitsPerSecond;

        this.syncLock = syncLock;
        this.redisTemplate = redisTemplate;
        this.maxBurstSeconds = maxBurstSeconds;
    }


    /**
     * 生成并存储默认令牌桶
     *
     * @return
     */
    private RedisPermits putDefaultPermits() {

        Lock lock = syncLock.obtain(lockKey);
        try {
            lock.lock();
            Object obj = redisTemplate.opsForValue().get(key);

            if (null == obj) {
                RedisPermits permits = new RedisPermits(permitsPerSecond, maxBurstSeconds);
                redisTemplate.opsForValue().set(key, permits, 1, TimeUnit.DAYS);
                return permits;
            } else {
                return (RedisPermits) obj;
            }
        } finally {
            lock.unlock();
        }


    }


    /**
     * 获取令牌桶
     *
     * @return
     */
    public RedisPermits getPermits() {
        try {
            Object obj = redisTemplate.opsForValue().get(key);
            if (null == obj) {
                return putDefaultPermits();
            }
            return (RedisPermits) obj;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 更新令牌桶
     *
     * @param permits
     */
    public void setPermits(RedisPermits permits) {
        redisTemplate.opsForValue().set(key, permits, 1, TimeUnit.DAYS);
    }


    /**
     * @param timeout 获取令牌等待的时间，负数被视为0
     * @param unit
     * @return
     */
    public Boolean tryAcquire(Long timeout, TimeUnit unit) {
        Long maxWaitingTime = System.currentTimeMillis() + unit.toMillis(timeout);
       // System.out.println("maxWaitingTime:"+maxWaitingTime);
        Lock lock = syncLock.obtain(lockKey);
        try {
            lock.lock();
            RedisPermits permits = this.getPermits();
            if (!permits.acquire(maxWaitingTime)) {
               return false;
            }
            this.setPermits(permits);
        } catch (Exception e) {
          //  e.printStackTrace();
        } finally {
            lock.unlock();
        }
        return true;
    }


}
